# -*- coding: utf-8 -*-

import os
from PyQt5 import QtCore
from PyQt5.QtGui import QImage, QPainter, QPen
from flow.label.stack import Stack


class LabeledBackendProcessor(object):
    def __init__(self, src_folder, dst_folder):
        self.src_stack = Stack(src_folder)
        self.dst_stack = Stack(dst_folder)
        self.current_node = None
        self.current_image = None
        self.show_image = None
        self.labeled_rectangles = list()
        self.left_top_point = None

    @staticmethod
    def read_labeled_rectangles(label_file_path):
        rectangles = []
        with open(label_file_path) as label_file:
            for line in label_file.readlines():
                line = line.strip().split()
                rectangle = [float(e) for e in line[1:]]
                # assert len(rectangle) == 4
                rectangles.append(rectangle)
        return rectangles

    def load_an_image(self):
        self.current_node = None
        self.current_image = None
        self.show_image = None
        self.labeled_rectangles = list()
        if self.src_stack.size() == 0:  # 所有图片已经处理完毕
            return self.dst_stack.size(), 0
        # 载入一张待处理的图片
        self.current_node = self.src_stack.pop()
        self.current_image = QImage(self.current_node.image_path)
        self.labeled_rectangles = self.read_labeled_rectangles(self.current_node.label_file_path)
        return self.dst_stack.size(), self.src_stack.size() + 1

    def next_image_button_event(self):
        if self.current_node is None:  # 当前没有图片在处理
            return self.dst_stack.size(), 0
        # step 1  保存当前的处理结果
        with open(self.current_node.label_file_path, "w") as label_file:
            for point in self.labeled_rectangles:
                str_line = str(0) + " " + " ".join([str(e) for e in point]) + "\n"
                label_file.writelines(str_line)
        # step 2  将图片和标注文件剪切到目标文件夹
        self.dst_stack.push(self.current_node)
        # step 3  载入下一张待处理的图片
        return self.load_an_image()

    def previous_image_button_event(self):
        # step 1  先保存当前的处理结果
        if self.current_node is not None:
            with open(self.current_node.label_file_path, "w") as label_file:
                for point in self.labeled_rectangles:
                    str_line = str(0) + " " + " ".join([str(e) for e in point]) + "\n"
                    label_file.writelines(str_line)
            self.src_stack.push(self.current_node)  # 将图片和标注文件保存回原文件夹
        # step 2  从目标文件夹提取回一张图片至源文件夹中
        if self.dst_stack.size() > 0:
            self.src_stack.push(self.dst_stack.pop())
        # step 3 从源文件夹中提取出一张图片进行处理
        return self.load_an_image()

    def delete_current_image(self):
        # step 1 删除当前的节点
        if self.current_node is not None:
            image_path = self.current_node.image_path
            label_file_path = self.current_node.label_file_path
            os.remove(image_path)
            os.remove(label_file_path)
        # step 2 从源文件夹中提取出一张图片进行处理
        return self.load_an_image()

    def set_left_top_point(self, relative_position):
        if self.current_image is None:
            return
        image_width = self.current_image.width()
        image_height = self.current_image.height()
        x = int(image_width * relative_position[0])
        y = int(image_height * relative_position[1])
        self.left_top_point = [x, y]

    def draw_image(self, relative_position=None, draw_dashed_line=False,
                   draw_mouse_move_rect=False, draw_mouse_release_rect=False):
        if self.current_image is None:
            return None
        assert isinstance(self.current_image, QImage)
        self.show_image = QImage(self.current_image)
        painter = QPainter()
        painter.begin(self.show_image)
        # 首先绘制所有已经标注的矩形框
        painter.setPen(QPen(QtCore.Qt.red, 3))
        for rectangle in self.labeled_rectangles:
            # assert len(rectangle) == 4
            image_width = self.current_image.width()
            image_height = self.current_image.height()
            c_x = int(rectangle[0] * image_width)
            c_y = int(rectangle[1] * image_height)
            w = int(rectangle[2] * image_width)
            h = int(rectangle[3] * image_height)
            x = c_x - w // 2
            y = c_y - h // 2
            painter.drawRect(x, y, w, h)
        # 绘制新的矩形
        if relative_position is not None:
            image_width = self.current_image.width()
            image_height = self.current_image.height()
            x = int(image_width * relative_position[0])
            y = int(image_height * relative_position[1])
            if draw_dashed_line:
                painter.setPen(QPen(QtCore.Qt.yellow, 2, QtCore.Qt.DashDotLine))
                painter.drawLine(0, y, image_width - 1, y)
                painter.drawLine(x, 0, x, image_height - 1)
            if draw_mouse_move_rect:
                painter.setPen(QPen(QtCore.Qt.blue, 3))
                if self.left_top_point is not None:
                    painter.drawRect(self.left_top_point[0],
                                     self.left_top_point[1],
                                     x - self.left_top_point[0] + 1,
                                     y - self.left_top_point[1] + 1)
            if draw_mouse_release_rect:
                if self.left_top_point is not None:
                    left_top_x = self.left_top_point[0]
                    left_top_y = self.left_top_point[1]
                    w = x - self.left_top_point[0] + 1
                    h = y - self.left_top_point[1] + 1
                    # 避免过小的矩形框出现，譬如鼠标误点击的情况
                    if w > 10 and h > 10:
                        painter.setPen(QPen(QtCore.Qt.red, 3))
                        if self.left_top_point is not None:
                            painter.drawRect(self.left_top_point[0], self.left_top_point[1], w, h)
                            # 保存成 yolo 需要的格式
                            center_x = left_top_x + w / 2
                            center_y = left_top_y + h / 2
                            self.labeled_rectangles.append([center_x / image_width,
                                                            center_y / image_height,
                                                            w / image_width,
                                                            h / image_height])
                    self.left_top_point = None
        painter.end()
        return self.show_image

    def right_mouse_button_event(self, x, y):
        # step 1  取消掉最后一次标注点
        if len(self.labeled_rectangles) > 0:
            for i, rect in enumerate(self.labeled_rectangles):
                if rect[0] - rect[2] / 2 <= x <= rect[0] + rect[2] / 2 and \
                        rect[1] - rect[3] / 2 <= y <= rect[1] + rect[3] / 2:
                    del self.labeled_rectangles[i]
                    break
            else:
                self.labeled_rectangles.pop()
        # step 2  在图像中重新绘制所有的标注点
        return self.draw_image()
